import aiohttp
import asyncio
import json
import re
import html
from data_v_1_0 import load_question_only
from tqdm import tqdm
from datetime import datetime
import os
"""
    v1.0:
    4/28
    多线程/协程
"""

# 配置
API_URL = "http://10.10.93.198:8090/v1/chat-messages"
API_KEY = "app-21QokC3toMRaHnbCzNhSKyqf"
# 测试qwen 32b app-21QokC3toMRaHnbCzNhSKyqf
# R1 app-SRoouEPrCLS0UW1uuwEThKhG
HEADERS = {
    "Authorization": f"Bearer {API_KEY}",
    "Content-Type": "application/json"
}
# sheet = "聊天助手"
OUTPUT_DIR = "../data"
os.makedirs(OUTPUT_DIR, exist_ok=True)


async def process_query(session, query, semaphore):
    """处理单个查询"""
    async with semaphore:  # 限制并发量
        payload = {
            "inputs": {},
            "query": query,
            "response_mode": "blocking",
            "conversation_id": "",
            "user": "abc-123",
            "files": []
        }

        try:
            async with session.post(
                    API_URL,
                    headers=HEADERS,
                    json=payload,
                    timeout=aiohttp.ClientTimeout(total=90) # TODO:90s就报超时err
            ) as response:

                response.raise_for_status()
                data = await response.json()
                answer = data.get("answer", "")
                # print(answer)

                # 1. 去除思维链： < details > 标签块（可选）
                # cleaned = re.sub(r"<details.*?</details>", "", answer, flags=re.DOTALL) TODO:这里的json格式为什么变了
                cleaned = re.sub(r"<think.*?</think>", "", answer, flags=re.DOTALL)

                # 2. 去除所有 HTML 标签
                cleaned = re.sub(r"<[^>]+>", "", cleaned)
                # 3. 将转义字符还原（如 \n, \", &nbsp; 等）
                cleaned_answer = html.unescape(cleaned)

                return {
                    "query": query,
                    "answer": cleaned_answer,
                    "error": None
                }

        except Exception as e:
            return {
                "query": query,
                "answer": f"[Error] {str(e)}",
                "error": str(e)
            }


async def save_results(results, filename):
    """实时保存结果到JSON"""
    with open(filename, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)


async def main(sheet="聊天助手"):
    # 加载测试集
    test_queries = load_question_only(sheet=sheet)
    results = []

    # 创建信号量控制并发数 (建议5-10之间) 10就崩了，唐老师建议最大到3吧
    semaphore = asyncio.Semaphore(3)

    # 进度条设置
    progress_bar = tqdm(total=len(test_queries), desc=f"处理 {sheet}")

    # 文件名带时间戳
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = f"{OUTPUT_DIR}/{sheet}_answers_{timestamp}.json"

    async with aiohttp.ClientSession() as session:
        tasks = []
        for query in test_queries:
            task = asyncio.create_task(process_query(session, query, semaphore))
            task.add_done_callback(lambda x: progress_bar.update(1))
            tasks.append(task)

        # 实时收集结果并保存
        for future in asyncio.as_completed(tasks):
            result = await future

            results.append(result)
            print(f"Query: {result['query'][:50]}... | Answer: {result['answer'][:50]}...")

            # 每10条保存一次
            if len(results) % 10 == 0:
                await save_results(results, output_file)

        # 最终保存
        await save_results(results, output_file)

    progress_bar.close()

    print(f"\n所有查询已完成，结果已保存到 {output_file}")
    print(f"成功: {len([r for r in results if not r['error']])} 失败: {len([r for r in results if r['error']])}")

    # print(output_file)
    return output_file # 返回结果文件地址，用于下一步评估


def run1(sheet="聊天助手"):
    output_file=asyncio.run(main(sheet="聊天助手"))
    return output_file

if __name__ == "__main__":
    asyncio.run(main())